import pylab as p
import scipy.signal as signal 

def filter_sinc(n, FPass, Fs):
    if n != 0:
        return p.sin(FPass * p.pi * n / Fs) / (p.pi * n)
    else:
        return FPass / float(Fs)

#------------------------------------------------
# Create a FIR filter
#------------------------------------------------
def createFIRFilter(Fs, FPass, N):
    
    sinc = [filter_sinc(n, (Fs - 2 * FPass), Fs) for n in p.arange(-121, 121 + 1, 1)]
    
    filter_coef = sinc * p.hanning(N)
    
    filter_coef = [filter_coef[n] * (-1) ** (n + 1) for n in range(len(filter_coef))]
    
    return filter_coef

#------------------------------------------------
# Create a IIR Butterworth filter
# http://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.iirfilter.html
#------------------------------------------------
def createIIRButterFilter(Fs, FPass, FStop, GPass, GStop):
    
    # The Nyquist rate of the signal.
    max_Freq = Fs / 2.0
#    
    N, Wn = signal.buttord(FPass / max_Freq, FStop / max_Freq, GPass, GStop, analog=0)

    b, a = signal.butter(N, Wn, btype='high', analog=0, output='ba')
    
    return [b, a]

#------------------------------------------------
# Respuesta al escalon
#------------------------------------------------    
def plotZeroesAndPoles(title, b, a=1):
    zeros = p.roots(b)
    poles = p.roots(a)
    p.title(title + ' - Polos y Ceros')
    if zeros.any():
        p.scatter(p.real(zeros), p.imag(zeros), c='b', marker='o')
    if poles.any():
        p.scatter(p.real(poles), p.imag(poles), c='r', marker='x')
    axes = p.axes()
    circle = p.Circle((0, 0), radius=1, alpha=1, facecolor="none")
    axes.add_patch(circle)  
    p.ylabel('Imaginario')
    p.xlabel('Real')
    p.axvline(x=0)
    p.axhline(y=0)
    
#------------------------------------------------
# Group delay del filtro
#------------------------------------------------    
def plotGroupDelay(title, b, a=1, Fs=1):
    w, h = signal.freqz(b, a)
    
    p.title(title + ' - Group Delay')
    delay = -p.diff(p.unwrap(p.arctan2(p.imag(h), p.real(h))))
    frecs = p.arange(0, float(Fs)/2, (float(Fs)/2) / len(delay))
    p.plot(frecs, delay)
    p.xlabel('Frecuencia(Hz)')
    p.ylabel('Delay(n)')
    
    
#------------------------------------------------
# Magnitud del filtro
#------------------------------------------------    
def plotMagnitudeResponse(title, b, a=1, Fs=1):
    w, h = signal.freqz(b, a)
    h_dB = 20 * signal.log10(abs(h))
    
    p.title(title + ' - Magnitud') 
    p.plot(w / max(w) * float(Fs) / 2, h_dB)
    p.ylabel('Magnitud(db)')
    p.xlabel('Frecuencia(Hz)')      
    
#------------------------------------------------
# Fase del filtro
#------------------------------------------------    
def plotPhaseResponse(title, b, a=1, Fs=1):
    w, h = signal.freqz(b, a)
    
    p.title(title + ' - Fase')
    h_Phase = p.unwrap(p.arctan2(p.imag(h), p.real(h)))
    p.plot(w / max(w) * float(Fs) / 2, h_Phase)
    p.ylabel('Fase(rad)')
    p.xlabel('Frecuencia(Hz)')    
    
#------------------------------------------------
# Respuesta al impulso
#------------------------------------------------    
def plotImpulseResponse(title, b, a=1):
    l = len(b)
    
    impulse = p.repeat(0., l); impulse[0] = 1.
    x = p.arange(0, l)
    response = signal.lfilter(b, a, impulse)
    
    p.subplot(211)
    p.title(title + ' - Respuesta al impulso')
    p.stem(x, response)
    p.ylabel('Amplitud')
    p.xlabel('n')
    
#------------------------------------------------
# Respuesta al escalon
#------------------------------------------------    
def plotStepResponse(title, b, a=1):
    l = len(b)
    
    impulse = p.repeat(0., l); impulse[0] = 1.
    x = p.arange(0, l)
    response = signal.lfilter(b, a, impulse)
          
    p.subplot(212)
    p.title(title + ' - Respuesta al escalon')
    step = p.cumsum(response)
    p.stem(x, step)
    p.ylabel('Amplitud')
    p.xlabel('n')
    p.subplots_adjust(hspace=0.5)    
